// ros
#include "ros/ros.h"

// bag sim
#include "vio_bag_sim.h"

namespace vio_bag_sim {

void VIOBagSim::Initialize() {
  LOG(INFO) << vio::LOG_GREEN_BEGIN << "load msgs from bag: " << bag_path_
            << vio::LOG_COLOR_END;
  LoadMsgsFromBag();
  LOG(INFO) << "init vio pipeline";
  vio_pipeline_.Initialize();
}

template <typename T>
void VIOBagSim::GetMsgsFromBag(const rosbag::Bag& bag,
                               const std::string& topic_name,
                               std::vector<T>* msgs_ptr) {
  rosbag::View view(bag, rosbag::TopicQuery(topic_name));
  for (auto& m : view) {
    msgs_ptr->push_back(*m.instantiate<T>());
  }
}

template <typename T>
void VIOBagSim::AddMsgsToMap(const MsgType& type,
                             const std::vector<T>& msgs,
                             std::multimap<double, MsgTypeIdx>* map_ptr) {
  MsgTypeIdx type_idx;
  type_idx.type = type;
  for (size_t i = 0; i < msgs.size(); i++) {
    type_idx.idx = i;
    const double ts = msgs[i].header.stamp.toSec();
    map_ptr->emplace(ts, type_idx);
  }
}

void VIOBagSim::LoadMsgsFromBag() {
  rosbag::Bag bag(bag_path_, rosbag::bagmode::Read);
  GetMsgsFromBag(bag, "/mynteye/imu/data_raw", &imus_);
  AddMsgsToMap(MsgType::IMU, imus_, &bag_msgs_map_);
  LOG(INFO) << "load IMU msgs: " << imus_.size();
  GetMsgsFromBag(bag, "/mynteye/left/image_raw", &left_imgs_);
  AddMsgsToMap(MsgType::LEFT_IMG, left_imgs_, &bag_msgs_map_);
  LOG(INFO) << "load left image msgs: " << left_imgs_.size();
  GetMsgsFromBag(bag, "/mynteye/right/image_raw", &right_imgs_);
  AddMsgsToMap(MsgType::RIGHT_IMG, right_imgs_, &bag_msgs_map_);
  LOG(INFO) << "load right image msgs: " << right_imgs_.size();
  bag.close();
}

void VIOBagSim::Step() {
  for (const auto& msg : bag_msgs_map_) {
    switch (msg.second.type) {
      case MsgType::IMU: {
        const sensor_msgs::Imu& imu = imus_.at(msg.second.idx);
        // add IMU
        const Eigen::Vector3d accel(imu.linear_acceleration.x,
                                    imu.linear_acceleration.y,
                                    imu.linear_acceleration.z);
        const Eigen::Vector3d gyro(imu.angular_velocity.x,
                                   imu.angular_velocity.y,
                                   imu.angular_velocity.z);
        vio_pipeline_.IMUCallback(imu.header.stamp.toSec(), accel, gyro);
        break;
      }
      case MsgType::LEFT_IMG: {
        const sensor_msgs::Image& img = left_imgs_.at(msg.second.idx);
        const cv_bridge::CvImagePtr cv_ptr =
            cv_bridge::toCvCopy(img, sensor_msgs::image_encodings::MONO8);
        // add left image
        vio_pipeline_.LeftImgCallback(img.header.stamp.toSec(), cv_ptr->image);
        break;
      }
      case MsgType::RIGHT_IMG: {
        const sensor_msgs::Image& img = right_imgs_.at(msg.second.idx);
        const cv_bridge::CvImagePtr cv_ptr =
            cv_bridge::toCvCopy(img, sensor_msgs::image_encodings::MONO8);
        vio_pipeline_.RightImgCallback(img.header.stamp.toSec(), cv_ptr->image);
        break;
      }
      default:
        break;
    }
  }
}

void VIOBagSim::Terminate() {
  vio_pipeline_.Terminate();
  LOG(INFO) << "bag sim finished";
}

}  // namespace vio_bag_sim
